#include "stdafx.h"
#include "../aclib/CBC.h"
#include "../aclib/Cryptor.h"
#include "../aclib/Block.h"
#include "../aclib/HexParser.h"
#include "MCryptor.h"

namespace AC = AppliedCryptography;
using ::testing::Return;
using ::testing::_;
using ::testing::Invoke;

template<unsigned int TBlockSize>
void SetBlock(const AC::Block<TBlockSize>& block1, AC::Block<TBlockSize>& block2)
{
	block2 = block1;
};

TEST(CBCChypertextStealing, EncryptLastBlockShorter)
{
	AC::ChypertextStealing<4, 5>::Encryption encCTS;
	MCryptor<4, 5> cryptor;
	AC::Block<5> secondToLastBlock;
	AC::Block<5> lastBlock;
	unsigned int lastBlockLength = 3;
	AC::HexParser Ox;
	AC::Block<5> chypertextN;
	AC::Block<5> chypertextNExp;

	// Step #3
	secondToLastBlock.Set(Ox("0102030405"));

	// Step #5
	lastBlock.Set(Ox("1020300000"));

	// Step #6
	EXPECT_CALL(cryptor, Process(_, _))
		.Times(1)
		.WillOnce(Invoke(SetBlock<5>));

	chypertextN = encCTS.Process(cryptor, lastBlock, lastBlockLength, secondToLastBlock);

	chypertextNExp.Set(Ox("0102030000"));
	EXPECT_TRUE(chypertextN == chypertextNExp);

	chypertextNExp.Set(Ox("1122330405"));
	EXPECT_TRUE(secondToLastBlock == chypertextNExp);
};

TEST(CBCChypertextStealing, EncryptLastBlockAlign)
{
	AC::ChypertextStealing<4, 5>::Encryption encCTS;
	MCryptor<4, 5> cryptor;
	AC::Block<5> secondToLastBlock;
	AC::Block<5> lastBlock;
	unsigned int lastBlockLength = 5;
	AC::HexParser Ox;
	AC::Block<5> chypertextN;
	AC::Block<5> chypertextNExp;

	// Step #3
	secondToLastBlock.Set(Ox("0102030405"));

	// Step #5
	lastBlock.Set(Ox("1020304050"));

	// Step #6
	EXPECT_CALL(cryptor, Process(_, _))
		.Times(1)
		.WillOnce(Invoke(SetBlock<5>));

	chypertextN = encCTS.Process(cryptor, lastBlock, lastBlockLength, secondToLastBlock);

	chypertextNExp.Set(Ox("0102030405"));
	EXPECT_TRUE(chypertextN == chypertextNExp);

	chypertextNExp.Set(Ox("1122334455"));
	EXPECT_TRUE(secondToLastBlock == chypertextNExp);
};

TEST(CBCChypertextStealing, DecryptLastBlockShorter)
{
	AC::ChypertextStealing<4, 5>::Decryption decCTS;
	MCryptor<4, 5> cryptor;
	AC::Block<5> secondToLastBlock;
	AC::Block<5> lastBlock;
	unsigned int lastBlockLength = 3;
	AC::HexParser Ox;
	AC::Block<5> chypertextN;
	AC::Block<5> chypertextNExp;

	// Step #3
	secondToLastBlock.Set(Ox("1122330405"));

	// Step #5
	lastBlock.Set(Ox("0102030000"));

	// Step #6
	EXPECT_CALL(cryptor, Process(_, _))
		.Times(1)
		.WillOnce(Invoke(SetBlock<5>));

	chypertextN = decCTS.Process(cryptor, lastBlock, lastBlockLength, secondToLastBlock);

	chypertextNExp.Set(Ox("1020300000"));
	EXPECT_TRUE(chypertextN == chypertextNExp);

	chypertextNExp.Set(Ox("0102030405"));
	EXPECT_TRUE(secondToLastBlock == chypertextNExp);
};

TEST(CBCChypertextStealing, DecryptLastBlockAlign)
{
	AC::ChypertextStealing<4, 5>::Decryption decCTS;
	MCryptor<4, 5> cryptor;
	AC::Block<5> secondToLastBlock;
	AC::Block<5> lastBlock;
	unsigned int lastBlockLength = 5;
	AC::HexParser Ox;
	AC::Block<5> chypertextN;
	AC::Block<5> chypertextNExp;

	// Step #3
	secondToLastBlock.Set(Ox("1122334455"));

	// Step #5
	lastBlock.Set(Ox("0102030405"));

	// Step #6
	EXPECT_CALL(cryptor, Process(_, _))
		.Times(1)
		.WillOnce(Invoke(SetBlock<5>));

	chypertextN = decCTS.Process(cryptor, lastBlock, lastBlockLength, secondToLastBlock);

	chypertextNExp.Set(Ox("1020304050"));
	EXPECT_TRUE(chypertextN == chypertextNExp);

	chypertextNExp.Set(Ox("0102030405"));
	EXPECT_TRUE(secondToLastBlock == chypertextNExp);
};
